import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler
from torch import optim
from torch.utils.data import DataLoader, Dataset
from config import Device, DataPath, batch_size, learning_rate, input_size, hidden_size, num_layers, epochs, \
    Val_Acc_Start, ModelPath, output_size
from data.loader import make_Dataset
from model.lstm_nn import new_LSTM
import time


def draw_figure(train_losses, train_accuracies, val_losses, val_accuracies):
    """绘制训练集和测试集的损失变化、准确率变化"""
    epochs = range(1, len(train_losses) + 1)

    # 绘制【损失】曲线
    plt.figure(figsize=(10, 6))
    plt.subplot(2, 1, 1)
    plt.plot(epochs, train_losses, label='Train Loss')
    plt.plot(epochs, val_losses, label='Validation Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()

    # 绘制【准确率】曲线
    plt.subplot(2, 1, 2)
    plt.plot(epochs, train_accuracies, label='Train Accuracy')
    plt.plot(epochs, val_accuracies, label='Validation Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()

    plt.tight_layout()
    plt.show()
    plt.savefig('training_metrics.png')


def train():
    # 记录每个epoch的损失和准确率
    train_losses = []
    train_accuracies = []
    val_losses = []
    val_accuracies = []

    # 初始化最佳验证损失
    best_val_acc = float('inf')  # 初始化最佳验证损失为﹢∞(正无穷)

    # 初始化模型、损失函数、优化器、学习率调整策略
    train_loader, val_loader = make_Dataset(DataPath, batch_size)
    model = new_LSTM(input_size, hidden_size, num_layers, output_size=output_size)
    model.to(Device)  # 将模型加载到Device上
    criterion = nn.MSELoss()  # 使用均方误差MSELoss作为损失函数
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # 使用Adam优化器，学习率设为0.001

    # 学习率调整策略
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10, verbose=True)
    scheduler_start = False

    for epoch in range(epochs):
        """训练模式"""
        model.train()  # 设置模型为训练模式
        teLossSum = 0.0
        teRightCnt = 0.0
        teSampleCnt = 0

        for inputs, labels in train_loader:
            inputs, labels = inputs.to(Device), labels.to(Device)  # 将数据加载到Device上
            optimizer.zero_grad()  # 清空梯度
            outputs = model(inputs)  # 前向传播
            loss = criterion(outputs, labels)  # 计算损失
            loss.backward()  # 反向传播
            optimizer.step()  # 更新权重

            # 批次损失和准确率
            teLossSum += loss.item()  # 累计损失
            _, predicted = torch.max(outputs, 1)  # 计算预测结果
            teRightCnt += (predicted == labels).sum().item()  # 累计预测正确的样本数
            teSampleCnt += labels.size(0)  # 累计样本数；量

        # 计算当前epoch下 训练集的平均损失和准确率
        train_losses.append(teLossSum / len(train_loader))
        train_accuracies.append(teRightCnt / teSampleCnt)

        """验证模式"""
        model.eval()  # 设置模型为验证模式
        veLossSum = 0.0
        veRightCnt = 0.0
        veSampleCnt = 0

        for inputs, labels in val_loader:
            inputs, labels = inputs.to(Device), labels.to(Device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            veLossSum += loss.item()

            # 批次损失和准确率
            _, predicted = torch.max(outputs, 1)
            veRightCnt += (predicted == labels).sum().item()
            veSampleCnt += labels.size(0)

        # 计算当前epoch下 验证集的平均损失和准确率
        val_losses.append(veLossSum / len(val_loader))
        val_accuracies.append(veRightCnt / veSampleCnt)

        # 输出当前epoch信息
        epoch_info = f"Epoch {(epoch + 1):5d}/{epochs}:, Train_Loss: {train_losses[-1]:.4f}, \tVal_Loss: {val_losses[-1]:.4f}, \tTrain_Acc: {train_accuracies[-1]:.4f}, \tVal_Acc: {val_accuracies[-1]:.4f}"
        print(epoch_info)

        """更新学习率"""
        if scheduler_start == False and val_accuracies[-1] >= Val_Acc_Start:
            scheduler_start = True
        if scheduler_start:
            scheduler.step(val_losses[-1])

        """模型保存"""
        if val_accuracies[-1] > best_val_acc and best_val_acc > 0.7:  # 验证集准确率大于0.7且大于当前最佳验证准确率时保存模型
            best_val_acc = val_accuracies[-1]
            model_path = ModelPath.replace('TIMESTAMP', str(int(time.time())) + "_" + epoch_info)
            torch.save(model.state_dict(), model_path)

    return train_losses, train_accuracies, val_losses, val_accuracies


if __name__ == '__main__':
    train_losses, train_accuracies, val_losses, val_accuracies = train()
    draw_figure(train_losses, train_accuracies, val_losses, val_accuracies)
